"""
智能体基类模块 - 定义智能体基类
"""

from typing import Dict, TypedDict, List, Union, Any
from abc import ABC, abstractmethod
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
import json
import re
import logging

from FDAS.config.settings import (
    OPENAI_API_KEY,
    OPENAI_BASE_URL,
    OPENAI_MODEL,
    OPENAI_TEMPERATURE
)

class BaseAgentState(TypedDict):
    """智能体状态基类"""
    user_query: str
    parsed_query: str
    analysis_result: str

class BaseAgent(ABC):
    """智能体基类"""
    
    def __init__(self, model: str = OPENAI_MODEL, temperature: float = OPENAI_TEMPERATURE):
        """
        初始化智能体
        
        参数:
            model: 大模型名称
            temperature: 温度参数
        """
        self.logger = logging.getLogger(self.__class__.__name__)
        
        # 初始化大模型
        self.llm = ChatOpenAI(
            model=model,
            base_url=OPENAI_BASE_URL,
            api_key=OPENAI_API_KEY,
            temperature=temperature
        )
        
        # 构建工作流
        self.workflow = self._build_workflow()
        
        # 编译执行图
        self.agent = self.workflow.compile()
    
    @abstractmethod
    def _build_workflow(self) -> StateGraph:
        """
        构建智能体工作流
        
        返回:
            工作流图
        """
        pass
    
    @abstractmethod
    def analyze(self, query: str) -> str:
        """
        执行分析
        
        参数:
            query: 用户查询
            
        返回:
            分析结果
        """
        pass
    
    def _clean_json_response(self, response: str) -> Dict:
        """
        清理并解析JSON响应
        
        参数:
            response: 大模型返回的JSON字符串
            
        返回:
            解析后的JSON字典
        """
        try:
            # 清理响应内容
            cleaned_content = re.sub(r'[\x00-\x1F\u200b-\u200d\ufe00-\ufe0f]', '', response)
            
            # 提取JSON - 使用Python支持的正则表达式
            # 尝试多种方法提取JSON，从最严格到最宽松
            # 1. 尝试找到整个响应中的第一个完整JSON对象
            json_match = None
            
            # 尝试使用json模块直接解析
            try:
                # 查找可能的JSON开始位置
                start_pos = cleaned_content.find('{')
                if start_pos != -1:
                    # 尝试不同的结束位置
                    for i in range(len(cleaned_content), start_pos, -1):
                        try:
                            potential_json = cleaned_content[start_pos:i]
                            json.loads(potential_json)
                            json_match = re.search(re.escape(potential_json), cleaned_content)
                            break
                        except json.JSONDecodeError:
                            continue
            except Exception:
                pass
            
            # 如果上面的方法失败，使用简单的正则表达式
            if not json_match:
                json_match = re.search(r'\{[^{}]*\}', cleaned_content)
            
            if not json_match:
                raise ValueError("未检测到有效JSON结构")
                
            # 解析JSON
            try:
                result = json.loads(json_match.group())
            except json.JSONDecodeError:
                # 尝试修复常见格式问题
                fixed_json = re.sub(r',\s*}', '}', json_match.group())
                fixed_json = re.sub(r',\s*]', ']', fixed_json)
                result = json.loads(fixed_json)
            
            return result
                
        except Exception as e:
            self.logger.error(f"JSON解析失败: {str(e)}")
            return {}
    
    def _format_error_response(self, error_message: str) -> str:
        """
        格式化错误响应
        
        参数:
            error_message: 错误信息
            
        返回:
            格式化后的错误响应
        """
        return f"""
## 分析过程中出现错误

很抱歉，在分析过程中遇到了以下错误：

```
{error_message}
```

请尝试以下解决方案：

1. 检查输入的资产代码是否正确
2. 确保选择了正确的资产类型
3. 尝试选择不同的分析类型或时间范围
4. 如果问题持续存在，请联系系统管理员
        """